package mao.redis.limitrate;

import mao.redis.jedispool.MyJedisPool;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.Pipeline;
import redis.clients.jedis.Response;

import java.io.IOException;
import java.util.concurrent.TimeUnit;

/**
 * @author Totoro
 * @date 2020/5/29
 * @explain:使用zset 做限流工作:使用当前时间作为socre每次请求时删除时间窗之外的值，然后zcard统计。
 */
public class SimpleRateLimit {
    Jedis jedis;
    public SimpleRateLimit(Jedis jedis) {
        this.jedis = jedis;
    }
    public boolean isAllowed(String userId, String userAction, int period, int maxCount) throws IOException{
        String key = String.format("srl:%s:%s", userId, userAction);
        Pipeline pipe = jedis.pipelined();
        pipe.multi();
        pipe.zadd(key,System.currentTimeMillis(),String.valueOf(System.currentTimeMillis()));
        //删除时间周期之外的。
        pipe.zremrangeByScore(key,0,System.currentTimeMillis()-period*1000);
        Response<Long> zcard = pipe.zcard(key);
        //设置过期时间，防止冷用户占用内存。
        pipe.expire(key,period+1);
        pipe.exec();
        pipe.syncAndReturnAll();
        pipe.close();
        Long aLong = zcard.get();
        return aLong<=maxCount;
    }

    public static void main(String[] args) throws Exception{
        Jedis jedis = MyJedisPool.getJedis();
        SimpleRateLimit rateLimit = new SimpleRateLimit(jedis);
        for (int i = 0; i < 20; i++) {
            //如果在一秒内有多次请求则会被记录为一次。
            TimeUnit.MICROSECONDS.sleep(500);
            System.out.println(rateLimit.isAllowed("Totoro","query",20,5));
        }
    }
}
